#!/usr/bin/env python3
# coding: utf-8
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# This file is a part of the CANN Open Software.
# Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
# Please refer to the License for details. You may not use this file except in compliance with the License.
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
# See LICENSE in the root of the software repository for the full text of the License.
# ======================================================================================================================
import random
import torch
import math 
import copy
from typing import Union, List

from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
from atk.case_generator.generator.base_generator import CaseGenerator
from atk.configs.case_config import InputCaseConfig, CaseConfig


# def scale_K_list(K_list, target_sum):
#     import math
#     K_sum = sum(K_list)
#     if K_sum <= target_sum:
#         return K_list
#     scale = target_sum / K_sum
#     K_temp = [max(1, math.floor(k * scale)) for k in K_list]
#     diff = target_sum - sum(K_temp)
#     # 按最大误差分配
#     errors = [(k * scale - math.floor(k * scale), idx) for idx, k in enumerate(K_list)]
#     if diff > 0:
#         errors.sort(reverse=True)
#         for i in range(diff):
#             K_temp[errors[i][1]] += 1
#     elif diff < 0:
#         errors.sort()
#         i = 0
#         while diff < 0 and i < len(K_list):
#             if K_temp[errors[i][1]] > 1:
#                 K_temp[errors[i][1]] -= 1
#                 diff += 1
#             i += 1
#     return K_temp

quantGroupSize = 256
K_opt = [256, 512, 768, 1024, 1280, 7168]
N_opt = [8, 16, 24,32,40, 48, 56, 64, 72,  256, 2800, 2808, 2816]

@GENERATOR_REGISTRY.register("ascend_generate_grouped_matmul_V4_a8w4")
class AscendGroupedMatmulV4A8W4(CaseGenerator):
    def __init__(self, config):
        super().__init__(config)

    def after_input_config(
            self,
            index: int,
            input_case: Union[InputCaseConfig, List[InputCaseConfig]]
    ) -> Union[InputCaseConfig, List[InputCaseConfig]]:

        return input_case

    def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
        x = case_config.inputs[0]
        weight = case_config.inputs[1]
        biasOptional = case_config.inputs[2]
        scaleOptional = case_config.inputs[3]
        perTokenScaleOptional = case_config.inputs[7]
        groupListOptional = case_config.inputs[8]
        print('x',case_config.inputs[0])
        print('weight',case_config.inputs[1])
        print('biasOptional',case_config.inputs[2])
        print('scaleOptional',case_config.inputs[3])
        print('perTokenScaleOptional',case_config.inputs[7])
        print('groupListOptional',case_config.inputs[8])

        print('======after_case_config前面================================')
        #print(case_config.inputs[0][0].shape[0])
        M = x[0].shape[0]
        K = random.choice(K_opt)
        x[0].shape[1] = K 
        E = weight[0].shape[0]
        if E > M :
            E = M
        weight[0].shape[0] = E
        weight[0].shape[1] = K
        N = random.choice(N_opt)
        weight[0].shape[2] = N
        biasOptional[0].shape[0],biasOptional[0].shape[1] = E,N 
        scaleOptional[0].shape[0],scaleOptional[0].shape[1],scaleOptional[0].shape[2] = E,int(K/quantGroupSize),N 
        perTokenScaleOptional[0].shape[0],perTokenScaleOptional[0].shape[1] = M,1
        groupListOptional.shape[0] = E

        print('======after_case_config之后================================')
        # print('x',case_config.inputs[0])
        # print('weight',case_config.inputs[1])
        # print('biasOptional',case_config.inputs[2])
        # print('scaleOptional',case_config.inputs[3])
        # print('perTokenScaleOptional',case_config.inputs[7])
        # print('groupListOptional',case_config.inputs[8])

        return case_config